#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
Catmull‑Rom 插值实现 4× 放大（128×128 → 512×512）
"""

from PIL import Image

def catmull_rom_interp(p0, p1, p2, p3, t, tau=0.5):
    """
    Catmull‑Rom 曲线插值（可调张力 τ）。
    p0‑p3 必须按顺序给出：左‑2、左‑1、当前、右‑1（或对应的垂直顺序）。
    """
    def to_tuple(v):
        return v if isinstance(v, tuple) else (v,)

    p0, p1, p2, p3 = map(to_tuple, (p0, p1, p2, p3))

    out = []
    for a0, a1, a2, a3 in zip(p0, p1, p2, p3):
        # 切向量（张力 τ）
        m1 = tau * (a2 - a0)
        m2 = tau * (a3 - a1)

        # Hermite 基函数
        h00 = 2 * t ** 3 - 3 * t ** 2 + 1
        h10 = t ** 3 - 2 * t ** 2 + t
        h01 = -2 * t ** 3 + 3 * t ** 2
        h11 = t ** 3 - t ** 2

        val = h00 * a1 + h10 * m1 + h01 * a2 + h11 * m2
        out.append(int(round(max(0, min(255, val)))))

    return out[0] if len(out) == 1 else tuple(out)


def get_pixel(img, x, y):
    """越界时返回最近的边缘像素（复制）"""
    w, h = img.size
    x = max(0, min(w - 1, x))
    y = max(0, min(h - 1, y))
    return img.getpixel((x, y))


def upscale_catmull_rom_4x(img, tau=0.5):
    """
    将输入图像放大四倍，使用 Catmull‑Rom 插值（可调张力 τ）。
    """
    src_w, src_h = img.size
    dst_w, dst_h = src_w * 4, src_h * 4
    out_img = Image.new(img.mode, (dst_w, dst_h))
    dst = out_img.load()

    # 对每个源像素生成 4×4 的目标块
    for y in range(src_h-1):
        for x in range(src_w-1):
            # ---------- 水平细分（5 条） ----------
            horiz = []
            for i in range(5):          # t = 0/4 … 4/4
                t = i / 4.0
                # 四个水平控制点：左‑2、左‑1、当前、右‑1
                p0 = get_pixel(img, x - 2, y)
                p1 = get_pixel(img, x - 1, y)
                p2 = get_pixel(img, x,     y)
                p3 = get_pixel(img, x + 1, y)
                horiz.append(catmull_rom_interp(p0, p1, p2, p3, t, tau))

            # ---------- 垂直细分（5 条） ----------
            vert = []
            for j in range(5):
                t = j / 4.0
                p0 = get_pixel(img, x, y - 2)
                p1 = get_pixel(img, x, y - 1)
                p2 = get_pixel(img, x, y)
                p3 = get_pixel(img, x, y + 1)
                vert.append(catmull_rom_interp(p0, p1, p2, p3, t, tau))

            # ---------- 交叉插值 ----------
            # 对每个 (i,j) 位置，先取水平细分点 i，再在垂直方向上对这 4 条水平结果做一次 Catmull‑Rom
            for i in range(5):
                for j in range(5):
                    # 这里我们把水平细分点视作“列”，再对列做垂直插值
                    # 取四个相邻的水平细分点作为控制点
                    ph0 = horiz[max(i - 2, 0)]
                    ph1 = horiz[max(i - 1, 0)]
                    ph2 = horiz[i]
                    ph3 = horiz[min(i + 1, 4)]

                    pixel = catmull_rom_interp(ph0, ph1, ph2, ph3, j / 4.0, tau)
                    dst[x * 4 + i, y * 4 + j] = pixel

    return out_img


if __name__ == "__main__":
    input_path  = "lena_128.png"
    output_path = "lena_512_catmull_rom.png"

    img = Image.open(input_path).convert("RGB")
    # τ = 0.5 为标准 Catmull‑Rom，0 为 B‑样条，1 为线性
    upscaled = upscale_catmull_rom_4x(img, tau=0.6)
    upscaled.save(output_path)
    print(f"已完成 4 倍放大，保存为 {output_path}")
